#include <iostream>
#include <vector>
#include <string>
#include <fstream>
#include <map>
#include <stdlib.h>
#include <lineages.h>
#include <algorithm>
#include <stdio.h>
#include <sstream>
#include <math.h>
#include <time.h>
#include <shared.h>

using namespace std;

int fitnessType;
int robustType;
int disturbType;
double baseDelRate = 1.0;

void verticalEvol2_mirror(const vector<int> & termTypes,vector<string> &resLineages,const vector<double> & depth2DivRate, int maxSamplePerGen, double fracSel);
// Sample from the possible divisions and use the good one.
// Sampling according to the division rate of that layer.

void verticalEvol3_mirror(const vector<int> & termTypes,vector<string> &resLineages, int maxSamplePerGen,double fracSel);
// Sample from the possible divisions and use the good one.
// Totally random sampling, do not follow the real depth distribution

void usage()
{
	cerr << "Perform macro-evolution, every generation would have a fix population size of 100." << endl;
    cerr << "This program is specifically designed for H.roretzi lineage," << endl;
    cerr << "where the two half-lineages are exactly the same, so that the random lineages retain this feature." << endl;
	cerr << "Usage: mirrorMacroEvol lin scheme fit robust disturb num seed collect frac" << endl;
	cerr << "Example: mirrorMacroEvol hro.alm 3 0 2 22 1 123 hro.macroEvol.lineages 0.95" << endl;
	cerr << "Note: All the lineages other than the H.roretzi lineage is handled differently, see \"macroEvol\" instead" << endl;
	cerr << "Parameters (all are needed) :"<< endl;
	cerr << " lin       : Name of a lineage file (one cell per line)" << endl;
	cerr << "             The macro-evolution are conducted based on this lineage." << endl;
    cerr << "             The file name must contain the string \"hro\" to indicate it is a H.roretzi lineage." << endl;
	cerr << "             The program need to know what lineage is being used because" << endl;
	cerr << "             it need to \"understand\" the names of functional classification" << endl;
	cerr << "             I recommend only those lineage files downloadable from the program's google code project site." << endl;
	cerr << " scheme    : Scheme of evolution." << endl;
	cerr << "             Could be 2 or 3" << endl;
	cerr << "             2. probability of adding one more division to a certain cell is determine by its depth," << endl;
	cerr << "                so that the resulting lineage would have similar terminal depths as the real lineage" << endl;
	cerr << "             3. every cell have same probability of adding one more division. (used in our paper)" << endl;
    cerr << " fit       : The type of \"fitness\" used during the evolution." << endl;
    cerr << "             Could be 0-4" << endl;
	cerr << "             0.  robustness" << endl;
	cerr << "             1.  (1-complexity)" << endl;
	cerr << "             2.  robustness/complexity" << endl;
	cerr << "             3.  robustness^5/complexity" << endl;
	cerr << "             4.  robustness^100/complexity" << endl;
	cerr << " robust    : If you used 0 or 2-4 for \"fit\", this parameter determines the type of robustness to be used." << endl;
	cerr << "             Could be 1-2" << endl;
	cerr << "             1.  product of fractional survival in each type, ignoring apoptotic cells (used in our paper)" << endl;
	cerr << "             2.  fractional survival of all non-apoptotic cells" << endl;
	cerr << " disturb   : Type of disturbance happen on the lineage. The robustness against this disturbance will be used." << endl;
	cerr << "             Could be 1 or 22" << endl;
	cerr << "             1.  a single necrosis among non-root cells" << endl;
	cerr << "             22.  program failure at the rate of 1/N_i, N_i is the number of internal nodes" << endl;
	cerr << "             Note that N_i is different at different stage of the macro-evolution." << endl;
	cerr << " num       : How many simulated lineages to evolve." << endl;
	cerr << " seed      : Seed for randomization" << endl;
	cerr << " collect   : Filename for the lineage collection (one lineage per line)" << endl;
	cerr << "             It must follow the same rule as \"lin\"" << endl;
	cerr << " frac      : In each generation of the evolution, a \"fraction\" of lineages with the highest fitness will be" << endl;
	cerr << "             chosen as the candidate progenitors of the next generation. The \"fraction\" here equals to (1-frac)." <<endl;
	cerr << "             The one progenitor of the next generation is randomly picked from these candidate progenitor lineages." << endl;
	cerr << " " << endl;
	return;
}

int main(int argc, char *argv[])
{
    if(argc != 10)
	{
		usage();
		exit(255);
	}
    string realLineagesFile = argv[1];
    int evolScheme = atoi(argv[2]);
    fitnessType = atoi(argv[3]);
    robustType = atoi(argv[4]);
    disturbType = atoi(argv[5]);
    int simLoop = atoi(argv[6]);
    int randSeed = atoi(argv[7]);
    string outputFile = argv[8];
    double fracSel = atof(argv[9]);
    srand(randSeed);

    int outputBuffer = 1;
    vector<vector<int> > typesPool;
    vector<vector<string> > lineagesPool;
    size_t found;

    found = realLineagesFile.find("mhr");
    if (found==string::npos)
    {
        cerr << "This program is only intended for the \"mhr\" (mirror hro) lineage." << endl;
        exit(0);
    }

    int popSize = 100;
    popSize = popSize;

    double realComplexity,realRobustness;
    vector<int> termTypes;
    vector<string> termLineages;

    loadRealLineages(realLineagesFile,termTypes,termLineages);
    Lineages realLin(termTypes,termLineages);

    realLin.getDepthStat();
    realComplexity = realLin.complexity();
    realLin.ruleDeletionRate = baseDelRate / (realLin.termTypes.size() - 1);
    // see misc.cpp for explanation for this
    realRobustness = realLin.robustness(disturbType,robustType);

    cout << "Parameters: FitnessType(" << fitnessType;
    cout << ") ; RobustType(" << robustType;
    cout << ") ; DistrubType (" << disturbType << ")";
    cout << endl;

    cout << "Lineage Name\tComplexity\tRobustness\tMax depth\tMean depth\tS.D. depth\tfitness" << endl;
    cout << realLineagesFile << "\t" << realComplexity << "\t" << realRobustness << "\t" << realLin.depthMax << "\t";
    cout << realLin.depthMean << "\t" << realLin.depthSD << "\t" << fitness(fitnessType,realRobustness,realComplexity) << endl;

    vector<int> halfTypes;
    vector<string> halfLineages;
    demirror(termLineages,termTypes,halfLineages,halfTypes);
    for(int i=0; i<simLoop; i++)
    {
        vector<int> randTypes;
        vector<string> randLineages;
        vector<string> randHalfLineages;

        random_shuffle(halfTypes.begin(),halfTypes.end());
        // shuffle the real to populate the random types

		if(evolScheme == 2)
		{
			verticalEvol2_mirror(halfTypes,randHalfLineages,realLin.depth2DivRate,100,fracSel);
			// mirror the lineage when calculating fitness or robustness
			// do NOT directly use realLin.depth2DivRate[i], use i +- 1, depending situation
		}
		else if(evolScheme == 3)
		{
			verticalEvol3_mirror(halfTypes,randHalfLineages,100,fracSel);
		}

        double simComplexity,simRobustness;
        mirror(randHalfLineages,halfTypes,randLineages,randTypes);

        Lineages simLin(randTypes,randLineages);
        simLin.getDepthStat();
        simComplexity = simLin.complexity();
        simLin.ruleDeletionRate = baseDelRate / (simLin.termTypes.size() - 1);
        simRobustness = simLin.robustness(disturbType,robustType);
        cout << i << "\t" << simComplexity << "\t" << simRobustness << "\t" << simLin.depthMax << "\t";
        cout << simLin.depthMean << "\t" << simLin.depthSD << "\t" << fitness(fitnessType,simRobustness,simComplexity)<< endl;

        typesPool.push_back(randTypes);
        lineagesPool.push_back(randLineages);
        if(i == 0)
        {
            writeLineageCollection(outputFile,typesPool,lineagesPool);
            typesPool.clear();
            lineagesPool.clear();
        }
        else if((i % outputBuffer == 0) or (i == simLoop - 1) )
        {
            appendLineageCollection(outputFile,typesPool,lineagesPool);
            typesPool.clear();
            lineagesPool.clear();
        }

    }
    return 0;
}


void verticalEvol2_mirror(const vector<int> & termTypes,vector<string> & resLineages,const vector<double> & depth2DivRate, int maxSamplePerGen,double fracSel)
{
    vector<int> curTypes; // current terminal types
    vector<vector<string> > pop; // the population of lineages (topology)
    vector<string> newParent;
    vector<string> indLin;
    vector<double> allFitness;
    vector<double> allFitnessSort;
    double fitnessCut;
    vector<double> allFitnessSE;
    vector<int>::iterator it;
//    double maxFitness;
    vector<int> maxFitnessIndex;
    int currentInd,randChild,otherChild,picked4Div;
    double tmpComplexity,tmpRobustness,rndDbl,sumDivRate;
    stringstream ss;
    string tmpStr;


    for(unsigned int i=0; i<termTypes.size(); i++) // each generation, will add one division to the lineage
    {
        picked4Div = 0;
        pop.clear();
        maxFitnessIndex.clear();
        allFitness.clear();
        allFitnessSE.clear();
        curTypes.push_back(termTypes[i]);
        if(newParent.size() == 0)
        {
            newParent.push_back("");
            allFitness.push_back(fitness(fitnessType,0.0,1.0));
            continue;
        }
        else
        {
            sumDivRate = 0.0;
            for(unsigned int j=0;j<newParent.size();j++)
            {
                sumDivRate += depth2DivRate[newParent[j].length()+1];
            }
            for(int j=0; j<maxSamplePerGen; j++)
            {
                rndDbl = random(0.0,sumDivRate);
                for(unsigned int k=0;k<newParent.size();k++)
                {
                    rndDbl -= depth2DivRate[newParent[k].length()+1];
                    if(rndDbl < 0.0)
                    {
                        picked4Div = k;
                        break;
                    }
                }

                pop.push_back(newParent);
                currentInd = pop.size() - 1;

                randChild = (int)random(0,2);
                otherChild = 1 - randChild;
                ss << pop[currentInd][picked4Div] << otherChild;
                tmpStr = ss.str();
                pop[currentInd].push_back(ss.str()); // the "new" daughter cell
                ss.str("");
                ss << pop[currentInd][picked4Div] << randChild;
                tmpStr = ss.str();
                pop[currentInd][picked4Div]= ss.str(); // the daughter cell whold holds her mother's property
                ss.str("");

                // construct lineage and calculate fitness
                vector<string> mirroredLins;
                vector<int> mirroredTypes;
                mirror(pop[currentInd],curTypes,mirroredLins,mirroredTypes);

                Lineages thisLin(mirroredTypes,mirroredLins);
                tmpComplexity = thisLin.complexity();
                thisLin.ruleDeletionRate = baseDelRate / (thisLin.termTypes.size() - 1 );
                // see misc.cpp for explanation of this
                tmpRobustness = thisLin.robustness(disturbType,robustType);
                allFitness.push_back(fitness(fitnessType,tmpRobustness,tmpComplexity));
            }

            allFitnessSort = allFitness;
            sort(allFitnessSort.begin(),allFitnessSort.end());
            fitnessCut = allFitnessSort[(int)(allFitnessSort.size() * fracSel)];

            for(unsigned int j = 0 ; j<allFitness.size(); j++)
            {
                if(allFitness[j] >= fitnessCut)
                {
                    maxFitnessIndex.push_back(j);
                }
            }

        }
        //cerr << "Division " << i << " : " << maxFitness << " A random lineage - R : " << tmpRobustness << " C : " << tmpComplexity << endl;
        rndDbl = random(0,maxFitnessIndex.size());
        newParent = pop[maxFitnessIndex[(int)rndDbl]];
    }

    resLineages = newParent;
    return;
}

void verticalEvol3_mirror(const vector<int> & termTypes,vector<string> & resLineages, int maxSamplePerGen, double fracSel)
{
    vector<int> curTypes; // current terminal types
    vector<vector<string> > pop; // the population of lineages (topology)
    vector<string> newParent;
    vector<string> indLin;
    vector<double> allFitness;
    vector<double> allFitnessSort;
    double fitnessCut;
    vector<double> allFitnessSE;
    vector<int>::iterator it;
//    double maxFitness;
    vector<int> maxFitnessIndex;
    int currentInd,randChild,otherChild,picked4Div;
    double tmpComplexity,tmpRobustness,rndDbl;
    stringstream ss;
    string tmpStr;


    for(unsigned int i=0; i<termTypes.size(); i++) // each generation, will add one division to the lineage
    {
        picked4Div = 0;
        pop.clear();
        maxFitnessIndex.clear();
        allFitness.clear();
        allFitnessSE.clear();
        curTypes.push_back(termTypes[i]);
        if(newParent.size() == 0)
        {
            newParent.push_back("");
            allFitness.push_back(fitness(fitnessType,0.0,1.0));
            continue;
        }
        else
        {
            for(int j=0; j<maxSamplePerGen; j++)
            {
                picked4Div = floor(random(0.0,newParent.size()));

                pop.push_back(newParent);
                currentInd = pop.size() - 1;

                randChild = (int)random(0,2);
                otherChild = 1 - randChild;
                ss << pop[currentInd][picked4Div] << otherChild;
                tmpStr = ss.str();
                pop[currentInd].push_back(ss.str()); // the "new" daughter cell
                ss.str("");
                ss << pop[currentInd][picked4Div] << randChild;
                tmpStr = ss.str();
                pop[currentInd][picked4Div]= ss.str(); // the daughter cell whold holds her mother's property
                ss.str("");

                // construct lineage and calculate fitness
                vector<string> mirroredLins;
                vector<int> mirroredTypes;
                mirror(pop[currentInd],curTypes,mirroredLins,mirroredTypes);

                Lineages thisLin(mirroredTypes,mirroredLins);
                tmpComplexity = thisLin.complexity();
                thisLin.ruleDeletionRate = baseDelRate / (thisLin.termTypes.size() - 1 );
                // see misc.cpp for explanation of this
                tmpRobustness = thisLin.robustness(disturbType,robustType);
                allFitness.push_back(fitness(fitnessType,tmpRobustness,tmpComplexity));
            }

            allFitnessSort = allFitness;
            sort(allFitnessSort.begin(),allFitnessSort.end());
            fitnessCut = allFitnessSort[(int)(allFitnessSort.size() * fracSel)];

            for(unsigned int j = 0 ; j<allFitness.size(); j++)
            {
                if(allFitness[j] >= fitnessCut)
                {
                    maxFitnessIndex.push_back(j);
                }
            }

        }
        //cerr << "Division " << i << " : " << maxFitness << " A random lineage - R : " << tmpRobustness << " C : " << tmpComplexity << endl;
        rndDbl = random(0,maxFitnessIndex.size());
        newParent = pop[maxFitnessIndex[(int)rndDbl]];
    }

    resLineages = newParent;
    return;
}
